#include "utilities.h"
#include "optimiz.h"
#include "lk.h"
#include "free.h"
#include "models.h"
#include "rates.h"
#include "asr.h"

#ifdef ENABLE_ASR

void Do_Asr(t_tree *tree)
{
	Prepare_Tree_For_Lk(tree);
	Lk(tree);

	if (tree->io->asr_algo == 1) // MARGINAL
	{
	  PhyML_Printf("\n. Writing marginal ancestral sequences in files '%s'_phyml_asr.node_N.dat,\n where N is a node number.\n", Basename(tree->io->in_align_file));
	  Marginal_Asr(tree->noeud[0], tree->noeud[0]->v[0], tree);
	}
	if (tree->io->asr_algo == 2) // JOINT
	{
	  PhyML_Printf("\n. Writing joint ancestral sequences. . .\n");
	  Make_Asr_Data(tree); // make data structures for the joint (=1) ASR algorithm.
	  Joint_Asr(tree->noeud[0], tree->noeud[0]->v[0], tree);
	  Free_Asr_Data(tree);
	}
}

// Make_Asr_Data builds data structures needed for joint ASR, but not marginal ASR.
void Make_Asr_Data(t_tree *tree)
{
	int i;
	int dim = tree->data->init_len;
	int nmixtures = 1;
#ifndef ENABLE_MBL
	nmixtures = tree->mod->n_catg;
#else
	nmixtures = n_l;
#endif
	For(i,2*tree->n_otu-1)
	{
		tree->noeud[i]->lxi = (phydbl *)mCalloc(nmixtures*dim*tree->mod->ns,sizeof(phydbl));
		tree->noeud[i]->cxi = (int *)mCalloc(nmixtures*dim*tree->mod->ns,sizeof(int));
		tree->noeud[i]->ppvector = (phydbl *)mCalloc(tree->mod->ns, sizeof(phydbl));
	}
}

// This method garbage collects the data structures allocated by Make_Asr_Data.
void Free_Asr_Data(t_tree *tree)
{
	int i;
	For(i,2*tree->n_otu-1)
	{
		Free(tree->noeud[i]->lxi);
		Free(tree->noeud[i]->cxi);
	}
}

// This method does marginal ASR for a single node d, whose parent is node a.
// state_pp is an array, where state_pp[state] = pp for this state
// ranked_states is used to order the PPs in descending order, for printing purposes.
void Marginal_Asr_Node(t_node *a, t_node *d, t_tree *tree, int site, phydbl *state_pp, int *ranked_states)
{
	int dim1 = tree->mod->n_catg * tree->mod->ns;
	int dim2 = tree->mod->ns;
	int mixture, state, i, is_left;

	// First,  sum the probs for this site ( = lk_sum)
	phydbl i_factor = 0.0;
	phydbl state_sum = 0.0;
	phydbl lk_sum = 0.0;

	For(state, tree->mod->ns)
	{
		state_sum = 1.0;
		For(i,3)
		{
			i_factor = 0.0;
			if (d == d->b[i]->left) is_left = 0;
			else is_left = 0;
#ifndef ENABLE_MBL
			For(mixture,tree->mod->n_catg)
#else
			For(mixture, n_l)
#endif
			{
#ifdef ENABLE_MBL
				if 		(is_left == 1 && !d->b[i]->left->tax) i_factor += d->b[i]->p_lk_left[site*dim1 + mixture*dim2 + state] * tree->mod->bl_mixture_wts[mixture];
				else if (is_left == 1 && d->b[i]->left->tax) i_factor += d->b[i]->p_lk_tip_l[site*dim2+state] * tree->mod->bl_mixture_wts[mixture];
				else if (is_left == 0 && !d->b[i]->rght->tax)  i_factor += d->b[i]->p_lk_rght[site*dim1 + mixture*dim2 + state] * tree->mod->bl_mixture_wts[mixture];
				else if (is_left == 0 && d->b[i]->rght->tax) i_factor += d->b[i]->p_lk_tip_r[site*dim2+state] * tree->mod->bl_mixture_wts[mixture];
#else
				if 		(is_left == 1 && !d->b[i]->left->tax) i_factor += d->b[i]->p_lk_left[site*dim1 + mixture*dim2 + state] * tree->mod->gamma_r_proba[mixture];
				else if (is_left == 1 && d->b[i]->left->tax) i_factor += d->b[i]->p_lk_tip_l[site*dim2+state] * tree->mod->gamma_r_proba[mixture];
				else if (is_left == 0 && !d->b[i]->rght->tax)  i_factor += d->b[i]->p_lk_rght[site*dim1 + mixture*dim2 + state] * tree->mod->gamma_r_proba[mixture];
				else if (is_left == 0 && d->b[i]->rght->tax) i_factor += d->b[i]->p_lk_tip_r[site*dim2+state] * tree->mod->gamma_r_proba[mixture];
#endif
			}

			// for debugging:
			int state2;
			For(state2, tree->mod->ns)
			{
				PhyML_Printf("\n. asr.c 97 P %d->%d, t= %f, Pij=%f", state, state2, d->b[i]->l, d->b[i]->Pij_rr[mixture*tree->mod->ns*tree->mod->ns + state*dim2 + state2]);
			}
			PhyML_Printf("\n. asr.c P(node %d branch %d site %d state %d) = %f", d->num, i, site, state, i_factor);
			// end debugging

			state_sum *= i_factor; // AND for each of 3 branches
		}
		PhyML_Printf("\n. asr.c P(node %d site %d state %d) = %f", d->num, site, state, state_sum);
		lk_sum += state_sum * tree->mod->pi[state]; // OR for each state
	}

	// Second, get PP for each state
	For(state, tree->mod->ns)
	{
		ranked_states[state] = state;
		state_pp[state] = 0.0;

		state_sum = 1.0;
		For(i,3)
		{
			if (d == d->b[i]->left) is_left = 1;
			else is_left = 0;
			i_factor = 0.0;
#ifdef ENABLE_MBL
			For(mixture,n_l)
#else
			For(mixture,tree->mod->n_catg)
#endif
			{
				phydbl this_part = 0.0;
#ifdef ENABLE_MBL
				if (is_left == 1) this_part += d->b[i]->p_lk_left[site*dim1 + mixture*dim2 + state] * tree->mod->bl_mixture_wts[mixture];
				else this_part += d->b[i]->p_lk_rght[site*dim1 + mixture*dim2 + state] * tree->mod->bl_mixture_wts[mixture];
#else
				if (is_left == 1) this_part += d->b[i]->p_lk_left[site*dim1 + mixture*dim2 + state] * tree->mod->gamma_r_proba[mixture];
				else this_part += d->b[i]->p_lk_rght[site*dim1 + mixture*dim2 + state] * tree->mod->gamma_r_proba[mixture];
#endif
				i_factor += this_part;
				//PhyML_Printf("\n. site %d state %d i %d mixture %d this_part %f", site, state, i, mixture, this_part);
			}
			state_sum *= i_factor;
		}
		//PhyML_Printf("\n. site %d state %d sum = %f", site, state, state_sum);
		state_pp[state] = state_sum * tree->mod->pi[state];
		state_pp[state] /= lk_sum;
	}
}

void Joint_Recur_Up_Terminal(t_node *x, t_tree *tree)
{
	PhyML_Printf("\n. asr.c 155 entered Joint_Recur_Up_Terminal x=%d", x->num, x->num);

	int dim = tree->data->init_len;
	int dim2 = tree->mod->ns;
	int dim3 = tree->mod->ns * tree->mod->ns;
	int site, a_state, x_state, mixture;
	For(site, dim)
	{
		//PhyML_Printf("asr.c 133 site = %d\n", site);
		For(a_state, tree->mod->ns)
		{
			x_state = tree->data->c_seq[x->num]->state[site];
			//PhyML_Printf("asr.c 127 x_state = %d\n", x_state);
			x_state = Assign_State_With_Ambiguity(&x_state,
								   tree->io->datatype,
								   tree->io->mod->state_len);
			//PhyML_Printf("asr.c 127 x_state = %d\n", x_state);

#ifdef ENABLE_MBL
			For(mixture,n_l)
			{	x->lxi[mixture*dim*dim2 + site*dim2 + a_state] = x->b[0]->Pij_rr[mixture*dim3+a_state*dim2+x_state];
#else
			For(mixture,tree->mod->n_catg)
			{	x->lxi[mixture*dim*dim2 + site*dim2 + a_state] = x->b[0]->Pij_rr[mixture*dim3+a_state*dim2+x_state];
#endif
			x->cxi[mixture*dim*dim2 + site*dim2 + a_state] = x_state;
			}
		}
	}
	return;
}

void Joint_Recur_Up_Internal(t_node *a, t_node *x, t_tree *tree)
{
	PhyML_Printf("\n. asr.c 155 entered Joint_Recur_Up_Internal a=%d d=%d", a->num, x->num);

	int dim = tree->data->init_len;
	int dim2 = tree->mod->ns;
	int dim3 = tree->mod->ns * tree->mod->ns;
	int site, a_state, x_state, mixture;
	int i, dir, a_dir;

	For(i, 3)
	{	if (x->b[i]->left == a || x->b[i]->rght == a)
		{	a_dir = i;
		}
	}

#ifdef ENABLE_MBL
	int nmixtures = n_l;
#else
	int nmixtures = tree->mod->n_catg;
#endif

	For(site, tree->data->init_len)
	{	For(a_state, tree->mod->ns)
		{
			// find the max_lxi and max_cxi for each a_state.
			phydbl *max_lxi = (phydbl *)mCalloc(nmixtures, sizeof(phydbl));
			int *max_cxi = (int *)mCalloc(nmixtures, sizeof(int));
			For(x_state, tree->mod->ns)
			{
#ifdef ENABLE_MBL
				For(mixture,n_l)
				{	phydbl this_lxi = x->b[a_dir]->Pij_rr[mixture*dim3+a_state*dim2+x_state];
#else
				For(mixture,tree->mod->n_catg)
				{	phydbl this_lxi = x->b[a_dir]->Pij_rr[mixture*dim3+a_state*dim2+x_state];
#endif
					//PhyML_Printf("\n. asr.c 175 %d %f %f", site, a_state, x_state);
					For(i, 3)
					{	if (x->v[i] != a)
						{	this_lxi *= x->v[i]->lxi[mixture*dim*dim2 + site*dim2 + x_state];
						}
					}
					//PhyML_Printf("\n. asr.c 194 this_lxi= %f, site %d a_state %d x_state %d mix %d", this_lxi, site, a_state, x_state, mixture);
					if (this_lxi > max_lxi[mixture])
					{
						max_lxi[mixture] = this_lxi;
						max_cxi[mixture] = x_state;
					}
				}
			}
#ifdef ENABLE_MBL
			For(mixture,n_l)
#else
			For(mixture,tree->mod->n_catg)
#endif
			{
				x->cxi[mixture*dim*dim2 + site*dim2 + a_state] = max_cxi[mixture];
				x->lxi[mixture*dim*dim2 + site*dim2 + a_state] = max_lxi[mixture];
			}
			Free(max_lxi);
			Free(max_cxi);
		}
	}
	return;
}

//
// a pre-order recursion of the tree, to calculate cxi and lxi values for every node.
//
void Joint_Recur_Up(t_node *a, t_node *d, t_tree *tree)
{
	int dim = tree->data->init_len;
	int dim2 = tree->mod->ns;
	int dim3 = tree->mod->ns * tree->mod->ns;
	int site, state, x_state, mixture;
	int i, dir;
	//PhyML_Printf("\n. asr.c 129 entering Joint_Recur_Up, node %d\n", d->num);
	if(d->tax) // d is a terminal taxon, so init lxi and cxi
	{
		Joint_Recur_Up_Terminal(d, tree);
	}
	else // Recur deeper until we hit taxa nodes
	{
		//PhyML_Printf("\n. asr.c 163 entering recursion loop\n");
		For(i,3)
		{
			//PhyML_Printf("\n. asr.c 169 i = %d\n", i);
			if (d->v[i] != a)
			{
				//PhyML_Printf("\n. asr.c 155 recurring to node node %d\n", d->v[i]->num);
				Joint_Recur_Up(d, d->v[i],tree);
			}
			else
			{	dir = i;
			}
		}
		if (d->v[dir]->tax) Joint_Recur_Up_Terminal(d->v[dir], tree);
		Joint_Recur_Up_Internal(d->v[dir], d, tree);
	}
}

void Joint_Recur_Down_Internal(t_node *a, t_node *x, t_tree *tree)
{
	// open asr write file
	char * out_asr_filepath = (char *)mCalloc(T_MAX_FILE,sizeof(char));
	Build_Asr_Filename(x, tree, out_asr_filepath);
	FILE * fout = Openfile(out_asr_filepath, 1);
	free(out_asr_filepath);

	int dim = tree->data->init_len;
	int dim2 = tree->mod->ns;
	int dim3 = tree->mod->ns * tree->mod->ns;
	int site, i, state, dir, mixture;

	// debugging:
	For(i, 3)
	{	if (x->v[i] != a)
		{
			PhyML_Printf("\n. asr.c 270 node %d is a child\n", x->v[i]->num);
		}
	}
	// end debugging:


	For(site, tree->data->init_len)
	{
		// goal: fill the ppvector for x.
		phydbl sum_pps = 0.0;
		For(state, tree->mod->ns)
		{
			phydbl r = tree->mod->pi[state];
			For(i, 3)
			{	if (x->v[i] != a)
				{
					phydbl mixture_sum = 0.0;
#ifdef ENABLE_MBL
					For(mixture,n_l)
					{	mixture_sum += x->v[i]->lxi[mixture*dim*dim2 + site*dim2 + state] * tree->mod->gamma_r_proba[mixture];
#else
					For(mixture,tree->mod->n_catg)
					{	mixture_sum += x->v[i]->lxi[mixture*dim*dim2 + site*dim2 + state] * tree->mod->gamma_r_proba[mixture];
#endif
					}
					PhyML_Printf("\n. asr.c 294 mixture_sum %f, site %d state %d node %d", mixture_sum, site, state, x->v[i]->num);
					r *= mixture_sum;
				}
			}
			x->ppvector[state] = r;
			sum_pps += r;
		}

		// Scale the lxi values into PPs
		For(state, tree->mod->ns)
		{	x->ppvector[state] = x->ppvector[state] / sum_pps;
		}

		// write for each state, using sum_lxi to get PPs
		// and then print
		PhyML_Fprintf(fout,"\n%d\t", site+1); fflush(fout);
		For(state, tree->mod->ns)
		{
			if (x->ppvector[state] > UNLIKELY_ANC_STATE) // only print ancestral states with non-zero PP support.
			{
				PhyML_Fprintf(fout,"%c %.3f\t", Reciproc_Assign_State(state, tree->io->datatype), x->ppvector[state]);
				fflush(fout);
			}
		}
	}

	// close the file
	fclose(fout);
}

void Joint_Recur_Down(t_node *a, t_node *x, t_tree *tree)
{
	// CASE: terminal node
	if(x->tax)
	{	return;
	}


	PhyML_Printf("\n. asr.c 322 writing ASR data for node %d from node %d\n", x->num, a->num);
	Joint_Recur_Down_Internal(a, x, tree);

	// recur to children nodes
	int i;
	For(i, 3)
	{	if (x->v[i] != a)
		{
		PhyML_Printf("\n. asr.c 330 recurring to node %d\n", x->v[i]->num);
		Joint_Recur_Down(x, x->v[i], tree);
		}
	}
}

//
// Does joint ASR, using an algorithm based the algorithm of Pupko et al., MBE 2000
//
void Joint_Asr(t_node *a, t_node *d, t_tree *tree)
{
	Joint_Recur_Up(a, d, tree);
	Joint_Recur_Down(a, d, tree);
}

// This method creates a filename for the ASR data for node d.
// out_asr_filepath is an empty string that has been allocated before calling this method,
// on completion of this method, out_asr_filepath will hold the filepath.
void Build_Asr_Filename(t_node *d, t_tree *tree, char *out_asr_filepath)
{
	// Determine output filename, and open the file for writing.
	char * nodename = (char *)mCalloc(T_MAX_LABEL,sizeof(char));
	sprintf(nodename, "%d", d->num + 1);
	//char * out_asr_filepath = (char *)mCalloc(T_MAX_FILE,sizeof(char));
	strcpy(out_asr_filepath,tree->io->in_align_file);
	strcat(out_asr_filepath,"_phyml_asr.node_");
	strcat(out_asr_filepath, nodename);
	if(tree->io->appebr_run_ID) { strcat(out_asr_filepath,"_"); strcat(out_asr_filepath,tree->io->run_id_string); }
	strcat(out_asr_filepath,".dat");
	free(nodename);
}

//
// Writes ancestral data to text files, for all nodes, recursively from node a.
// a = ancestral node
// d = descendant node
//
void Marginal_Asr(t_node *a, t_node *d, t_tree *tree)
{
	// first, recur.
	int i;
	if(d->tax)
    {	return;
    }
	else
    {
      For(i,3)
      {	if(d->v[i] != a) Marginal_Asr(d, d->v[i], tree);
      }
    }

	// Open the file (fout) for writing the ancestral states.
	char * out_asr_filepath = (char *)mCalloc(T_MAX_FILE,sizeof(char));
	Build_Asr_Filename(d, tree, out_asr_filepath);
	FILE * fout = Openfile(out_asr_filepath, 1);
	Free(out_asr_filepath);

	// Do ASR on every site.
	int site, state;
	phydbl *state_pp = (phydbl *)mCalloc(tree->mod->ns,sizeof(phydbl));
	int *ranked_states = (int *)mCalloc(tree->mod->ns,sizeof(phydbl));
	For(site, tree->data->init_len)
	{
		// First, get PPs for states
		Marginal_Asr_Node(a, d, tree, tree->data->sitepatt[site], state_pp, ranked_states);

		// sort the scores in descending order
		Quick_sort(0, tree->mod->ns - 1, state_pp, ranked_states);

		// and then print the ancestral states
		PhyML_Fprintf(fout,"\n%d\t", site+1); fflush(fout);
		For(state, tree->mod->ns)
		{
			if (state_pp[state] > UNLIKELY_ANC_STATE) // only print ancestral states with non-zero PP support.
			{
				PhyML_Fprintf(fout,"%c %.3f\t", Reciproc_Assign_State(ranked_states[state], tree->io->datatype), state_pp[state]);
				fflush(fout);
			}
		}
	}
	// cleanup:
	PhyML_Fprintf(fout, "\n");
	Free(state_pp);
	Free(ranked_states);
	fclose(fout);
}

// Partition is used in Quicksort
int Partition(int low, int high, phydbl pps[], int states[])
{
	// debugging code:
	//PhyML_Printf("\n. asr.c 448 entered Partition, low=%d, high=%d\n", low, high);
	//int qqq;
	//for(qqq=low; qqq<=high; qqq++)
	//{
	//	PhyML_Printf("%d(%f) ", states[qqq], pps[qqq]);
	//}
	//PhyML_Printf("\n");

	phydbl x = pps[low];
	int i = low - 1;
	int j = high + 1;
	phydbl temp_d;
	int temp_i;
	do
	{
	   do
	   {	j --;
	   }while (x > pps[j]);

	   do
	   {	i++;
	   }while (x < pps[i]);

	   if (i < j)
	   {
			 temp_d = pps[i];
			 pps[i] = pps[j];
			 pps[j] = temp_d;

			 temp_i = states[i];
			 states[i] = states[j];
			 states[j] = temp_i;
	   }
	}while (i < j);
	return j;           // returns middle subscript
}

void Quick_sort(int low, int high, phydbl pps[], int states[])
{
    int middle;
    if (low < high)
    {
        middle = Partition(low, high, pps, states);
        Quick_sort(low, middle, pps, states);   // sort first section
        Quick_sort(middle+1, high, pps, states);    // sort second section
    }
    return;
}

#endif
